#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>
#include <iostream>

struct Detection {
    float x1, y1, x2, y2, score;
    int class_id;
};

cv::Mat preprocess_image(cv::Mat& image) {
    cv::Mat float_img;
    cv::cvtColor(image, image, cv::COLOR_BGR2RGB);
    image.convertTo(float_img, CV_32F, 1.0 / 255.0);

    // 转换为CHW格式
    std::vector<cv::Mat> channels(3);
    cv::split(float_img, channels);
    cv::Mat chw;
    cv::merge(channels, chw); // 合并为CHW
    return chw;
}

Ort::Value create_input_tensor(Ort::Session& session, cv::Mat& processed_image) {
    // 获取输入形状信息
    Ort::TypeInfo input_type = session.GetInputTypeInfo(0);
    // auto input_shape = input_type.GetTensorTypeAndShapeInfo().GetShape();
    std::vector<int64_t> input_shape = {1,3,processed_image.rows,processed_image.cols};
    
    // 创建输入张量
    size_t input_size = 1;
    for (auto dim : input_shape) input_size *= dim; // 计算总元素数
    
    std::vector<float> input_data(input_size);
    memcpy(input_data.data(), processed_image.data, processed_image.total() * sizeof(float));
    
    return Ort::Value::CreateTensor<float>(
        Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault),
        input_data.data(), input_data.size(),
        input_shape.data(), input_shape.size()
    );
}

std::vector<Ort::Value> run_inference(
    Ort::Session& session, 
    Ort::Value& input_tensor,
    const char* input_name,
    const char* output_name
) {
    const char* input_names[] = {input_name};
    const char* output_names[] = {output_name};
    
    return session.Run(
        Ort::RunOptions{nullptr}, 
        input_names, &input_tensor, 1, 
        output_names, 1
    );
}


std::vector<Detection> parse_output(Ort::Value& output_tensor, const cv::Size& orig_size,cv::Size size_input, float conf_thresh=0.5) {
    float* data = output_tensor.GetTensorMutableData<float>();
    auto shape = output_tensor.GetTensorTypeAndShapeInfo().GetShape();
    
    float scale_x = orig_size.width / (float)size_input.width;
    float scale_y = orig_size.height / (float)size_input.height;
    
    std::vector<Detection> detections;
    for (int i = 0; i < shape[1]; ++i) {
        float* box = data + i * shape[2];
        if (box[4] < conf_thresh) continue;
        
        Detection det;
        det.x1 = box[0] * scale_x;
        det.y1 = box[1] * scale_y;
        det.x2 = box[2] * scale_x;
        det.y2 = box[3] * scale_y;
        det.score = box[4];
        det.class_id = static_cast<int>(box[5]);
        detections.push_back(det);
    }
    
    // 实现NMS函数，我不会写，先等一等
    // nms(detections, 0.45);
    
    for (const auto& det : detections) {
        
        std::cout << "Class: " << det.class_id << " | Conf: " << det.score 
                  << " | Box: [" << det.x1 << "," << det.y1 << "," 
                  << det.x2 << "," << det.y2 << "]\n";
    }
    return detections;
}

void draw_detections(cv::Mat& image, const std::vector<Detection>& detections) {
    for (const auto& det : detections) {
        // 绘制矩形框
        cv::rectangle(image, 
            cv::Point(det.x1, det.y1), 
            cv::Point(det.x2, det.y2), 
            cv::Scalar(0, 255, 0),  // 绿色框
            2  // 线宽
        );

        // 构建标签文本（类别 + 置信度）
        std::string label = "Class " + std::to_string(det.class_id) 
                           + ": " + std::to_string(det.score).substr(0, 4);

        // 在框的左上角添加标签
        cv::putText(image, 
            label, 
            cv::Point(det.x1, det.y1 - 5), 
            cv::FONT_HERSHEY_SIMPLEX, 
            0.5,  // 字体大小
            cv::Scalar(0, 255, 0),  // 绿色文本
            1  // 线宽
        );
    }
}

int main() {
    Ort::Env env;
    Ort::Session session(env, "/home/hitcrt/crt/test/demo/data/best.onnx", Ort::SessionOptions{});
    //注意/home/hitcrt/crt/test/demo/data/best.onnx是绝对路径，记得换！！！！
    Ort::AllocatorWithDefaultOptions allocator;

    cv::Mat image;
    image = cv::imread("/home/hitcrt/crt/test/demo/data/1.png");
    // cv::Size target_size(640, 640);//这个我想要书中指
    cv::Mat resized_image;
    cv::resize(image, resized_image, cv::Size(640, 640));
    cv::Mat processed = preprocess_image(resized_image);
    
    // 创建输入张量
    Ort::Value input_tensor = create_input_tensor(session, processed);
    Ort::AllocatedStringPtr input_name_ptr = session.GetInputNameAllocated(0, allocator);
    const char* input_name = input_name_ptr.get();
    Ort::AllocatedStringPtr output_name_ptr = session.GetOutputNameAllocated(0, allocator);
    const char* output_name = output_name_ptr.get(); 

    auto outputs = session.Run(Ort::RunOptions{nullptr}, 
                                 &input_name, &input_tensor, 1, 
                                 &output_name, 1);
        
    std::vector<Detection> detections = parse_output(outputs[0], image.size(),cv::Size(640,640));
    
    cv::Mat image_processed = image.clone();
    draw_detections(image_processed,detections);
    cv::imshow("original",image);
    cv::imshow("processed",image_processed);
    cv::waitKey(0);
    return 0;
}